%{
#include "ast.h"
#include <memory>
#include <vector>
#include <cstdio>

extern int yylex();
extern FILE* yyin;
void yyerror(const char* s);

CompUnit* g_comp_unit = nullptr;
%}

%union {
    int num;
    char* str;
    Expr* expr;
    Stmt* stmt;
    Param* param;
    FuncDef* func_def;
    std::vector<Stmt*>* stmt_list;
    std::vector<Expr*>* expr_list;
    std::vector<Param*>* param_list;
    std::vector<FuncDef*>* func_list;
}

%token <num> NUMBER
%token <str> ID
%token INT VOID IF ELSE WHILE BREAK CONTINUE RETURN
%token LPAREN RPAREN LBRACE RBRACE SEMI COMMA END
%token ASSIGN EQ NEQ LT LE GT GE AND OR
%token NOT PLUS MINUS MUL DIV MOD

%type <expr> Expr LorExpr LAndExpr EqExpr RelExpr AddExpr MulExpr UnaryExpr PrimaryExpr
%type <stmt> Stmt Block
%type <param> Param
%type <func_def> FuncDef
%type <stmt_list> stmt_list
%type <expr_list> expr_list
%type <param_list> param_list
%type <func_list> func_list

%start CompUnit

%%

CompUnit:
    func_list { g_comp_unit = new CompUnit(*$1); }
    ;

stmt_list:
    /* empty */ { $$ = new std::vector<Stmt*>; }
    | Stmt { $$ = new std::vector<Stmt*>; $$->push_back($1); }
    | stmt_list Stmt { $1->push_back($2); $$ = $1; }
    ;

Stmt:
    Block { $$ = $1; }
    | SEMI { $$ = new SBlock(std::vector<Stmt*>{}); }
    | Expr SEMI { $$ = new SExpr($1); }
    | ID ASSIGN Expr SEMI { $$ = new SAssign(std::string($1), $3);  }
    | INT ID ASSIGN Expr SEMI { $$ = new SDecl(std::string($2), $4); }
    | IF LPAREN Expr RPAREN Stmt ELSE Stmt { $$ = new SIf($3, $5, $7); }
    | IF LPAREN Expr RPAREN Stmt { $$ = new SIf($3, $5); }
    | WHILE LPAREN Expr RPAREN Stmt { $$ = new SWhile($3, $5); }
    | BREAK SEMI { $$ = new SBreak(); }
    | CONTINUE SEMI { $$ = new SContinue(); }
    | RETURN Expr SEMI { $$ = new SReturn($2); }
    | RETURN SEMI { $$ = new SReturn(); }
    ;

Block:
    LBRACE stmt_list RBRACE { $$ = new SBlock(*$2); }
    ;

func_list:
    FuncDef { $$ = new std::vector<FuncDef*>; $$->push_back($1); }
    | func_list FuncDef { $1->push_back($2); $$ = $1; }
    ;

FuncDef:
    INT ID LPAREN param_list RPAREN Block {
        $$ = new FuncDef{
             FuncType::TInt, 
             std::string($2), 
             *$4, 
             std::unique_ptr<Stmt>($6)
         };
         free($2);
    }
    | VOID ID LPAREN param_list RPAREN Block {
        $$ = new FuncDef{
             FuncType::TVoid, 
             std::string($2), 
             *$4, 
             std::unique_ptr<Stmt>($6)
         };
         free($2);
    }
    ;

param_list:
    /* empty */ { $$ = new std::vector<Param*>; }
    | Param { $$ = new std::vector<Param*>; $$->push_back($1); }
    | param_list COMMA Param { $1->push_back($3); $$ = $1; }
    ;

Param:
    INT ID { $$ = new Param{std::string($2)}; }
    ;

expr_list:
    /* empty */ { $$ = new std::vector<Expr*>; }
    | Expr { $$ = new std::vector<Expr*>; $$->push_back($1); }
    | expr_list COMMA Expr { $1->push_back($3); $$ = $1; }
    ;

Expr:
    LorExpr { $$ = $1; }
    ;

LorExpr:
    LAndExpr { $$ = $1; }
    | LorExpr OR LAndExpr { $$ = new EBinOp(BinOp::Or, $1, $3); }
    ;

LAndExpr:
    EqExpr { $$ = $1; }
    | LAndExpr AND EqExpr { $$ = new EBinOp(BinOp::And, $1, $3); }
    ;

EqExpr:
    RelExpr { $$ = $1; }
    | EqExpr EQ RelExpr { $$ = new EBinOp(BinOp::Eq, $1, $3); }
    | EqExpr NEQ RelExpr { $$ = new EBinOp(BinOp::Neq, $1, $3); }
    ;

RelExpr:
    AddExpr { $$ = $1; }
    | RelExpr LT AddExpr { $$ = new EBinOp(BinOp::Lt, $1, $3); }
    | RelExpr GT AddExpr { $$ = new EBinOp(BinOp::Gt, $1, $3); }
    | RelExpr LE AddExpr { $$ = new EBinOp(BinOp::Le, $1, $3); }
    | RelExpr GE AddExpr { $$ = new EBinOp(BinOp::Ge, $1, $3); }
    ;

AddExpr:
    MulExpr { $$ = $1; }
    | AddExpr PLUS MulExpr { $$ = new EBinOp(BinOp::Add, $1, $3); }
    | AddExpr MINUS MulExpr { $$ = new EBinOp(BinOp::Sub, $1, $3); }
    ;

MulExpr:
    UnaryExpr { $$ = $1; }
    | MulExpr MUL UnaryExpr { $$ = new EBinOp(BinOp::Mul, $1, $3); }
    | MulExpr DIV UnaryExpr { $$ = new EBinOp(BinOp::Div, $1, $3); }
    | MulExpr MOD UnaryExpr { $$ = new EBinOp(BinOp::Mod, $1, $3); }
    ;

UnaryExpr:
    PrimaryExpr { $$ = $1; }
    | PLUS UnaryExpr %prec NOT { $$ = $2; }
    | MINUS UnaryExpr %prec NOT { $$ = new EUnOp(UnOp::Neg, $2); }
    | NOT UnaryExpr { $$ = new EUnOp(UnOp::Not, $2); }
    ;

PrimaryExpr:
    ID { $$ = new EId(std::string($1)); }
    | NUMBER { $$ = new EInt($1); }
    | LPAREN Expr RPAREN { $$ = $2; }
    | ID LPAREN expr_list RPAREN { $$ = new ECall(std::string($1), *$3); }
    ;

%%

void yyerror(const char* s) {
    fprintf(stderr, "语法错误: %s\n", s);
}
